#include "CUDA_Poisson.h"
#include "utility.h"
#include <stdio.h>
#include "cuPrintf.cu"

static const int BLOCK_SIZE = 16;
static const int R = 1;


__global__ void poisson_kernel(const float * x, float * Ax, int dim_size)
{
	int ix = blockIdx.x * blockDim.x + threadIdx.x;
	int iy = blockIdx.y * blockDim.y + threadIdx.y;
	if( ix < dim_size && iy < dim_size )
	{
		float value = 4 * x[iy * dim_size + ix];
		value += (-1) * (ix > 0) * x[max(iy * dim_size + (ix-1),0)];
		value += (-1) * (iy > 0) * x[max((iy-1) * dim_size + ix,0)];
		value += (-1) * (ix < dim_size-1) * x[min(iy * dim_size + (ix + 1), dim_size * dim_size - 1)];
		value += (-1) * (iy < dim_size-1) * x[min((iy+1) * dim_size + ix, dim_size * dim_size - 1)];
		Ax[iy * dim_size + ix] = value;
	}
}

__global__ void poisson_kernel_shared_mem(const float * x, float * Ax, int dim_size)
{
	__shared__ float vec_cache [BLOCK_SIZE + 2*R][BLOCK_SIZE + 2*R];
	int ix = blockIdx.x * blockDim.x + threadIdx.x;
	int iy = blockIdx.y * blockDim.y + threadIdx.y;
	if( ix < dim_size && iy < dim_size )
	{
		//load values into cache
		int cache_x = threadIdx.x + R; int cache_y = threadIdx.y + R;
		vec_cache[cache_x][cache_y] = x[iy * dim_size + ix];
		if(threadIdx.x==0)
			vec_cache[cache_x-1][cache_y] = (ix > 0 ) * x[max(iy * dim_size + ix-1, 0)];
		else if((threadIdx.x==blockDim.x - 1) || (ix==dim_size-1))
			vec_cache[cache_x+1][cache_y] = (ix < dim_size-1) * x[min(iy * dim_size + ix+1,dim_size * dim_size -1)];
		if(threadIdx.y==0)
			vec_cache[cache_x][cache_y - 1] = (iy > 0 ) * x[max((iy-1) * dim_size + ix, 0)];
		else if((threadIdx.y==blockDim.y - 1) || (iy==dim_size-1))
			vec_cache[cache_x][cache_y+1] = (iy < dim_size-1) * x[min((iy+1) * dim_size + ix, dim_size * dim_size -1)];

		__syncthreads();

		//compute with cache
		float value = 4 * vec_cache[cache_x][cache_y];
		value -= vec_cache[cache_x-1][cache_y];
		value -= vec_cache[cache_x+1][cache_y];
		value -= vec_cache[cache_x][cache_y-1];
		value -= vec_cache[cache_x][cache_y+1];
		Ax[iy * dim_size + ix] = value;

	}
}

void CUDA_Poisson::operator()(const CUDA_vec<float> & a, CUDA_vec<float> & b) const
{
	int vec_dim_size = (int)sqrt(a.getSize());
	int num_blocks_dim = (int)ceil(((float)vec_dim_size) / BLOCK_SIZE);
	dim3 blockDim(BLOCK_SIZE,BLOCK_SIZE);
	dim3 gridDim(num_blocks_dim, num_blocks_dim);
//	cudaPrintfInit(24*1024*1024);
	if(use_shmem)
		poisson_kernel_shared_mem<<<gridDim, blockDim>>>(a.raw_pointer(),b.raw_pointer(), vec_dim_size);
	else
		poisson_kernel<<<gridDim, blockDim>>>(a.raw_pointer(),b.raw_pointer(), vec_dim_size);
//	cudaPrintfDisplay(stdout, true);
//	cudaPrintfEnd();
//	print_CUDA_err_msg("Poisson Multiply");
}

void CUDA_Poisson::getDiagonal(CUDA_vec<float> & a) const
{
	a.set_constant(4);
}
